package info.wangyuan.agent.service.curd.impl;

import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.google.common.collect.Sets;
import info.wangyuan.agent.common.Enum.Action;
import info.wangyuan.agent.entity.dto.AddKnowledgeSourceDTO;
import info.wangyuan.agent.entity.po.KnowledgeBase;
import info.wangyuan.agent.entity.po.KnowledgeSource;
import info.wangyuan.agent.mapper.KnowledgeSourceMapper;
import info.wangyuan.agent.service.curd.CommonService;
import info.wangyuan.agent.service.curd.KnowledgeSourceService;
import info.wangyuan.agent.service.event.KbAndAgentBindsChangedEvent;
import org.springframework.context.ApplicationEventPublisher;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import java.util.List;
import java.util.Objects;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * 知识库来源服务实现类
 *
 * @author Albert
 * @since 2025-07-14
 */
@Service
public class KnowledgeSourceServiceImpl extends ServiceImpl<KnowledgeSourceMapper, KnowledgeSource> implements KnowledgeSourceService {

    private final ApplicationEventPublisher publisher;
    private final CommonService commonService;

    public KnowledgeSourceServiceImpl(ApplicationEventPublisher publisher, CommonService commonService) {
        this.publisher = publisher;
        this.commonService = commonService;
    }

    /**
     * 批量添加或更新知识源
     */
    @Override
    @Transactional
    public void addOrUpdateByBatch(AddKnowledgeSourceDTO addKsDTO) {
        Integer agentId = addKsDTO.getAgentId();
        Set<Integer> newKbIds = addKsDTO.getKnowledgeBaseIds() != null
                ? Set.copyOf(addKsDTO.getKnowledgeBaseIds())
                : Set.of();

        // 获取当前 Agent 的已有知识库关联，并保证 null 安全
        Set<Integer> oldKbIds = list(Wrappers.lambdaQuery(KnowledgeSource.class)
                .eq(KnowledgeSource::getAgentId, agentId))
                .stream()
                .map(KnowledgeSource::getKnowledgeBaseId)
                .filter(Objects::nonNull)
                .collect(Collectors.toSet());

        // 计算新增和删除的差集
        Set<Integer> addKbIds = Sets.difference(newKbIds, oldKbIds);
        Set<Integer> removeKbIds = Sets.difference(oldKbIds, newKbIds);

        // 统一处理删除和新增
        if (!addKbIds.isEmpty())
            processKbChanges(agentId, addKbIds, Action.ADD);
        if (!removeKbIds.isEmpty())
            processKbChanges(agentId, removeKbIds, Action.REMOVE);
    }

    /**
     * 处理知识库变更并发布事件
     */
    private void processKbChanges(Integer agentId, Set<Integer> kbIds, Action action) {
        if (kbIds.isEmpty()) return;

        // 删除或新增数据库关联
        if (action == Action.REMOVE) {
            remove(Wrappers.lambdaQuery(KnowledgeSource.class)
                    .eq(KnowledgeSource::getAgentId, agentId)
                    .in(KnowledgeSource::getKnowledgeBaseId, kbIds));
        } else if (action == Action.ADD) {
            Set<KnowledgeSource> newSources = kbIds.stream()
                    .map(kbId -> KnowledgeSource.builder()
                            .agentId(agentId)
                            .knowledgeBaseId(kbId)
                            .build())
                    .collect(Collectors.toSet());
            if (!newSources.isEmpty())
                saveBatch(newSources);
        }

        // 获取 KnowledgeBase 对象并发布事件
        List<KnowledgeBase> kbs = commonService.getKbs(kbIds);
        if (!kbs.isEmpty())
            publisher.publishEvent(new KbAndAgentBindsChangedEvent(agentId, kbs, action));
    }
}
